# -*- coding: UTF-8 -*-
'''
@File ：5_bert_cls.py
@IDE ：PyCharm
@Author ：chaojie
@Date ：2025/11/7 
@Introduce:  基于bert 进行分类
'''
import json
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from transformers.models.bert import BertTokenizer, BertModel

bert_id=r'D:\linuxFiles\nlp_demo\weights\bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(bert_id)

class cls_dataset(Dataset):
    def __init__(self, data_path, label2id_file):
        super().__init__()

        self.labe2id = self.get_label2id(label2id_file)
        self.datas = self.load_data(data_path)


    def load_data(self, data_file):
        data = []
        data_list = pd.read_csv(data_file, header=None, sep='\t')

        for words, label_name in zip(data_list[0], data_list[1]):
            label = self.labe2id[label_name]    # 转换为 id
            data.append({"word": words, 'label': label})

        return data

    def get_label2id(self, label2id_file):
        df = pd.read_csv(label2id_file, sep=',')
        label2id = {}

        for name, label in zip(df['cls_name'], df['id']):
            label2id[name] = label
        return label2id

    def __len__(self):
        return len(self.datas)

    def __getitem__(self, index):
        return self.datas[index]


def collate_fn(batch):





    texts = [dir_['word'] for dir_ in batch]
    label = [dir_['label'] for dir_ in batch]

    texts = tokenizer(
        texts,
        truncation=True,
        padding='max_length',
        max_length=512,
        return_tensors='pt'
    )


    return texts, torch.tensor(label)


if __name__ == '__main__':
    label2if_file   = r'D:\linuxFiles\nlp_demo\datas/cls/label_2_id.csv'
    data_path       = r'D:\linuxFiles\nlp_demo\datas/cls/val.csv'

    dataset = cls_dataset(data_path, label2if_file)
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=0, collate_fn=collate_fn)
    for batch in dataloader:
        print(batch)